#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
姿态分类模块
作者: 元歌
开发地: 广州南沙铁建中心
时间: 2024年10月
"""

import numpy as np
import time
from typing import List, Tuple, Optional, Dict
from pose_detector import PoseDetector

class PoseClassifier:
    """姿态分类器，用于判断人体姿态（站立/坐下/倒下）"""
    
    def __init__(self, 
                 use_advanced_features: bool = True,
                 smoothing_window: int = 8,
                 confidence_threshold: float = 0.7):
        """
        初始化姿态分类器
        
        Args:
            use_advanced_features: 是否使用高级特征
            smoothing_window: 平滑窗口大小
            confidence_threshold: 置信度阈值
        """
        self.detector = PoseDetector()
        self.use_advanced_features = use_advanced_features
        self.smoothing_window = smoothing_window
        self.confidence_threshold = confidence_threshold
        
        # 姿态历史用于平滑
        self.pose_history = []
        self.max_history = 15
        self.min_pose_duration = 2.0  # 最小姿态持续时间（秒）
        self.last_pose_change = time.time()
        
    def classify_pose(self, landmarks) -> Dict[str, any]:
        """
        分类人体姿态
        
        Args:
            landmarks: 姿态关键点对象
            
        Returns:
            包含分类结果的字典
        """
        if not landmarks:
            return {
                'pose': 'unknown',
                'confidence': 0.0,
                'features': {},
                'reason': 'No landmarks detected'
            }
        
        # 提取关键特征
        if self.use_advanced_features:
            features = self._extract_advanced_features(landmarks)
        else:
            features = self._extract_features(landmarks)
        
        # 分类姿态
        pose, confidence, reason = self._classify(features)
        
        # 应用平滑处理
        if self.smoothing_window > 1:
            pose, confidence = self._apply_smoothing(pose, confidence)
        
        return {
            'pose': pose,
            'confidence': confidence,
            'features': features,
            'reason': reason
        }
    
    def _extract_features(self, landmarks) -> Dict[str, float]:
        """提取用于分类的关键特征"""
        features = {}
        
        # 添加运动检测相关特征
        features['movement_detected'] = False
        features['step_frequency'] = 0.0
        
        # 获取关键点坐标
        left_hip = self.detector.get_landmark_by_name(landmarks, 'left_hip')
        right_hip = self.detector.get_landmark_by_name(landmarks, 'right_hip')
        left_knee = self.detector.get_landmark_by_name(landmarks, 'left_knee')
        right_knee = self.detector.get_landmark_by_name(landmarks, 'right_knee')
        left_ankle = self.detector.get_landmark_by_name(landmarks, 'left_ankle')
        right_ankle = self.detector.get_landmark_by_name(landmarks, 'right_ankle')
        left_shoulder = self.detector.get_landmark_by_name(landmarks, 'left_shoulder')
        right_shoulder = self.detector.get_landmark_by_name(landmarks, 'right_shoulder')
        nose = self.detector.get_landmark_by_name(landmarks, 'nose')
        
        if not all([left_hip, right_hip, left_knee, right_knee, left_ankle, right_ankle, left_shoulder, right_shoulder, nose]):
            return features
        
        # 计算髋部中心点
        hip_center = (
            (left_hip[0] + right_hip[0]) / 2,
            (left_hip[1] + right_hip[1]) / 2,
            (left_hip[2] + right_hip[2]) / 2
        )
        
        # 计算肩部中心点
        shoulder_center = (
            (left_shoulder[0] + right_shoulder[0]) / 2,
            (left_shoulder[1] + right_shoulder[1]) / 2,
            (left_shoulder[2] + right_shoulder[2]) / 2
        )
        
        # 计算膝盖中心点
        knee_center = (
            (left_knee[0] + right_knee[0]) / 2,
            (left_knee[1] + right_knee[1]) / 2,
            (left_knee[2] + right_knee[2]) / 2
        )
        
        # 计算脚踝中心点
        ankle_center = (
            (left_ankle[0] + right_ankle[0]) / 2,
            (left_ankle[1] + right_ankle[1]) / 2,
            (left_ankle[2] + right_ankle[2]) / 2
        )
        
        # 1. 身体高度（肩部到髋部的距离）
        body_height = self.detector.calculate_distance(shoulder_center, hip_center)
        features['body_height'] = body_height
        
        # 2. 腿部长度（髋部到脚踝的距离）
        leg_length = self.detector.calculate_distance(hip_center, ankle_center)
        features['leg_length'] = leg_length
        
        # 3. 大腿长度（髋部到膝盖的距离）
        thigh_length = self.detector.calculate_distance(hip_center, knee_center)
        features['thigh_length'] = thigh_length
        
        # 4. 小腿长度（膝盖到脚踝的距离）
        shin_length = self.detector.calculate_distance(knee_center, ankle_center)
        features['shin_length'] = shin_length
        
        # 5. 膝盖角度（髋部-膝盖-脚踝的角度）
        left_knee_angle = self.detector.calculate_angle(left_hip, left_knee, left_ankle)
        right_knee_angle = self.detector.calculate_angle(right_hip, right_knee, right_ankle)
        features['left_knee_angle'] = left_knee_angle
        features['right_knee_angle'] = right_knee_angle
        features['avg_knee_angle'] = (left_knee_angle + right_knee_angle) / 2
        
        # 6. 身体倾斜角度（肩部-髋部连线与垂直线的夹角）
        body_angle = self._calculate_body_angle(shoulder_center, hip_center)
        features['body_angle'] = body_angle
        
        # 7. 头部相对位置（鼻子相对于髋部的位置）
        head_relative_y = nose[1] - hip_center[1]
        features['head_relative_y'] = head_relative_y
        
        # 8. 膝盖相对位置（膝盖相对于髋部的位置）
        knee_relative_y = knee_center[1] - hip_center[1]
        features['knee_relative_y'] = knee_relative_y
        
        # 9. 脚踝相对位置（脚踝相对于髋部的位置）
        ankle_relative_y = ankle_center[1] - hip_center[1]
        features['ankle_relative_y'] = ankle_relative_y
        
        # 10. 身体比例（身体高度与腿部长度的比例）
        if leg_length > 0:
            features['body_leg_ratio'] = body_height / leg_length
        else:
            features['body_leg_ratio'] = 0
        
        return features
    
    def _calculate_body_angle(self, shoulder_center: Tuple[float, float, float], 
                             hip_center: Tuple[float, float, float]) -> float:
        """计算身体倾斜角度"""
        # 计算身体向量
        body_vector = np.array([shoulder_center[0] - hip_center[0], shoulder_center[1] - hip_center[1]])
        
        # 垂直向量
        vertical_vector = np.array([0, -1])
        
        # 计算角度
        cos_angle = np.dot(body_vector, vertical_vector) / (np.linalg.norm(body_vector) * np.linalg.norm(vertical_vector))
        cos_angle = np.clip(cos_angle, -1.0, 1.0)
        angle = np.arccos(cos_angle)
        
        return np.degrees(angle)
    
    def _classify(self, features: Dict[str, float]) -> Tuple[str, float, str]:
        """
        基于特征分类姿态
        
        Args:
            features: 提取的特征字典
            
        Returns:
            (姿态类型, 置信度, 原因)
        """
        if not features:
            return 'unknown', 0.0, 'No features available'
        
        # 获取关键特征
        avg_knee_angle = features.get('avg_knee_angle', 180)
        body_angle = features.get('body_angle', 0)
        head_relative_y = features.get('head_relative_y', 0)
        knee_relative_y = features.get('knee_relative_y', 0)
        ankle_relative_y = features.get('ankle_relative_y', 0)
        body_leg_ratio = features.get('body_leg_ratio', 1)
        
        # 分类逻辑
        confidence = 0.0
        reason = ""
        
        # 检查是否倒下（身体倾斜角度过大）
        if body_angle > 45:
            confidence = min(0.9, body_angle / 90)
            reason = f"Body angle too large: {body_angle:.1f}°"
            return 'lying_down', confidence, reason
        
        # 检查是否坐下（膝盖角度较小，膝盖位置较高）
        if avg_knee_angle < 120 and knee_relative_y > -0.1:
            confidence = 0.8
            reason = f"Knee angle small: {avg_knee_angle:.1f}°, knee position high"
            return 'sitting', confidence, reason
        
        # 检查是否站立（膝盖角度较大，身体相对直立）
        if avg_knee_angle > 140 and body_angle < 30:
            confidence = 0.8
            reason = f"Knee angle large: {avg_knee_angle:.1f}°, body angle small: {body_angle:.1f}°"
            return 'standing', confidence, reason
        
        # 检查是否半蹲（膝盖角度中等）
        if 120 <= avg_knee_angle <= 140:
            confidence = 0.6
            reason = f"Medium knee angle: {avg_knee_angle:.1f}°"
            return 'squatting', confidence, reason
        
        # 默认情况
        confidence = 0.3
        reason = f"Unclear pose - knee_angle: {avg_knee_angle:.1f}°, body_angle: {body_angle:.1f}°"
        return 'unknown', confidence, reason
    
    def get_pose_description(self, pose: str) -> str:
        """获取姿态的中文描述"""
        descriptions = {
            'standing': '站立',
            'sitting': '坐下',
            'lying_down': '倒下/躺下',
            'squatting': '蹲下',
            'unknown': '未知姿态'
        }
        return descriptions.get(pose, '未知姿态')
    
    def _extract_advanced_features(self, landmarks) -> Dict[str, float]:
        """提取高级特征用于更精确的分类"""
        features = self._extract_features(landmarks)
        
        # 添加更多高级特征
        try:
            # 身体倾斜角度
            features['body_tilt_angle'] = self._calculate_body_tilt(landmarks)
            
            # 重心位置
            features['center_of_mass'] = self._calculate_center_of_mass(landmarks)
            
            # 肢体对称性
            features['limb_symmetry'] = self._calculate_limb_symmetry(landmarks)
            
            # 运动速度（如果有历史数据）
            features['motion_velocity'] = self._calculate_motion_velocity(landmarks)
            
            # 姿态稳定性
            features['pose_stability'] = self._calculate_pose_stability(landmarks)
            
        except Exception as e:
            print(f"高级特征提取出错: {e}")
        
        return features
    
    def _calculate_body_tilt(self, landmarks) -> float:
        """计算身体倾斜角度"""
        try:
            # 使用肩膀和臀部计算身体倾斜
            left_shoulder = self.detector.get_landmark_by_name(landmarks, 'left_shoulder')
            right_shoulder = self.detector.get_landmark_by_name(landmarks, 'right_shoulder')
            left_hip = self.detector.get_landmark_by_name(landmarks, 'left_hip')
            right_hip = self.detector.get_landmark_by_name(landmarks, 'right_hip')
            
            if all([left_shoulder, right_shoulder, left_hip, right_hip]):
                # 计算肩膀和臀部的中心点
                shoulder_center = ((left_shoulder[0] + right_shoulder[0]) / 2, 
                                 (left_shoulder[1] + right_shoulder[1]) / 2)
                hip_center = ((left_hip[0] + right_hip[0]) / 2, 
                             (left_hip[1] + right_hip[1]) / 2)
                
                # 计算倾斜角度
                dx = shoulder_center[0] - hip_center[0]
                dy = shoulder_center[1] - hip_center[1]
                angle = np.arctan2(dy, dx) * 180 / np.pi
                
                return angle
        except:
            pass
        return 0.0
    
    def _calculate_center_of_mass(self, landmarks) -> Tuple[float, float]:
        """计算重心位置"""
        try:
            # 使用主要关键点计算重心
            key_points = ['nose', 'left_shoulder', 'right_shoulder', 
                         'left_hip', 'right_hip', 'left_knee', 'right_knee']
            
            points = []
            for point_name in key_points:
                point = self.detector.get_landmark_by_name(landmarks, point_name)
                if point:
                    points.append(point)
            
            if points:
                x_coords = [p[0] for p in points]
                y_coords = [p[1] for p in points]
                return (np.mean(x_coords), np.mean(y_coords))
        except:
            pass
        return (0.5, 0.5)
    
    def _calculate_limb_symmetry(self, landmarks) -> float:
        """计算肢体对称性"""
        try:
            # 比较左右肢体的角度
            left_arm_angle = self._calculate_arm_angle(landmarks, 'left')
            right_arm_angle = self._calculate_arm_angle(landmarks, 'right')
            left_leg_angle = self._calculate_leg_angle(landmarks, 'left')
            right_leg_angle = self._calculate_leg_angle(landmarks, 'right')
            
            arm_symmetry = 1.0 - abs(left_arm_angle - right_arm_angle) / 180.0
            leg_symmetry = 1.0 - abs(left_leg_angle - right_leg_angle) / 180.0
            
            return (arm_symmetry + leg_symmetry) / 2.0
        except:
            pass
        return 1.0
    
    def _calculate_arm_angle(self, landmarks, side: str) -> float:
        """计算手臂角度"""
        try:
            shoulder = self.detector.get_landmark_by_name(landmarks, f'{side}_shoulder')
            elbow = self.detector.get_landmark_by_name(landmarks, f'{side}_elbow')
            wrist = self.detector.get_landmark_by_name(landmarks, f'{side}_wrist')
            
            if all([shoulder, elbow, wrist]):
                # 计算手臂角度
                v1 = np.array([elbow[0] - shoulder[0], elbow[1] - shoulder[1]])
                v2 = np.array([wrist[0] - elbow[0], wrist[1] - elbow[1]])
                
                cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
                angle = np.arccos(np.clip(cos_angle, -1.0, 1.0)) * 180 / np.pi
                return angle
        except:
            pass
        return 90.0
    
    def _calculate_leg_angle(self, landmarks, side: str) -> float:
        """计算腿部角度"""
        try:
            hip = self.detector.get_landmark_by_name(landmarks, f'{side}_hip')
            knee = self.detector.get_landmark_by_name(landmarks, f'{side}_knee')
            ankle = self.detector.get_landmark_by_name(landmarks, f'{side}_ankle')
            
            if all([hip, knee, ankle]):
                # 计算腿部角度
                v1 = np.array([knee[0] - hip[0], knee[1] - hip[1]])
                v2 = np.array([ankle[0] - knee[0], ankle[1] - knee[1]])
                
                cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
                angle = np.arccos(np.clip(cos_angle, -1.0, 1.0)) * 180 / np.pi
                return angle
        except:
            pass
        return 90.0
    
    def _calculate_motion_velocity(self, landmarks) -> float:
        """计算运动速度"""
        # 这里需要历史数据，暂时返回0
        return 0.0
    
    def _calculate_pose_stability(self, landmarks) -> float:
        """计算姿态稳定性"""
        # 基于关键点的变化程度计算稳定性
        return 1.0
    
    def _apply_smoothing(self, pose: str, confidence: float) -> Tuple[str, float]:
        """应用平滑处理减少抖动"""
        # 添加当前姿态到历史
        self.pose_history.append((pose, confidence))
        if len(self.pose_history) > self.max_history:
            self.pose_history = self.pose_history[-self.max_history:]
        
        if len(self.pose_history) < self.smoothing_window:
            return pose, confidence
        
        # 使用最近几帧进行平滑
        recent_poses = self.pose_history[-self.smoothing_window:]
        
        # 统计最频繁的姿态
        pose_counts = {}
        total_confidence = 0
        
        for p, c in recent_poses:
            pose_counts[p] = pose_counts.get(p, 0) + 1
            total_confidence += c
        
        # 选择最频繁的姿态
        most_common_pose = max(pose_counts, key=pose_counts.get)
        avg_confidence = total_confidence / len(recent_poses)
        
        # 检查是否需要保持当前姿态（基于时间）
        current_time = time.time()
        time_since_change = current_time - self.last_pose_change
        
        # 如果最频繁姿态的占比足够高，且满足时间要求，使用它
        if (pose_counts[most_common_pose] / len(recent_poses) >= 0.7 and 
            time_since_change >= self.min_pose_duration):
            self.last_pose_change = current_time
            return most_common_pose, avg_confidence
        else:
            # 保持当前姿态，但更新置信度
            return pose, confidence


